package cn.smileyan.demo;

import org.apache.flink.api.java.utils.MultipleParameterTool;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch7.ElasticsearchSink;
import org.apache.http.HttpHost;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 创建索引 hello-world 并向这个索引中写入数据
 * @author Smileyan
 */
public class FlinkElasticsearchDemo {
    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        final MultipleParameterTool parameters = MultipleParameterTool.fromArgs(args);

        // 设置 Elasticsearch 集群的连接信息 比如 http://localhost:9200 | http://localhost:9200,http://localhost:9200
        String urls = parameters.get("es.hosts", "http://localhost:9200");
        final String regex = ",";
        List<HttpHost> httpHosts = Arrays.stream(urls.split(regex)).map(HttpHost::create).collect(Collectors.toList());

        // 创建 ElasticsearchSinkFunction 用于将数据写入 Elasticsearch
        final String index = parameters.get("index", "hello-world");
        ElasticsearchSink.Builder<String> esSinkBuilder = buildSinkEs(httpHosts, index);

        // 创建数据流
        DataStream<String> resultStream = env.fromElements("data1", "data2", "data3");

        // 将结果 Sink 到 Elasticsearch
        resultStream.addSink(esSinkBuilder.build());

        // 执行 Flink 作业
        env.execute("Flink Elasticsearch Example");
    }

    /**
     * 获取用于将字符串数据 Sink 到 Elasticsearch 的 ElasticsearchSink.Builder 对象。
     *
     * @param httpHosts Elasticsearch 集群的连接信息
     * @param index 存储到 Elasticsearch 中的索引
     * @return ElasticsearchSink.Builder 对象
     */
    private static ElasticsearchSink.Builder<String> buildSinkEs(List<HttpHost> httpHosts, String index) {
        ElasticsearchSinkFunction<String> elasticsearchSinkFunction = (element, ctx, indexer) -> {
            // 将数据写入 Elasticsearch
            Map<String, String> json = new HashMap<>(1);
            json.put("data", element);
            IndexRequest source = Requests.indexRequest().index(index).source(json);
            indexer.add(source);
        };

        // 创建 ElasticsearchSink
        ElasticsearchSink.Builder<String> esSinkBuilder = new ElasticsearchSink.Builder<>(httpHosts, elasticsearchSinkFunction);

        // 设置批量写入的缓冲区大小，可根据实际情况调整
        esSinkBuilder.setBulkFlushMaxActions(1);
        return esSinkBuilder;
    }

}
